// --------------------------------------------------------------------------------------------------------------------
// <copyright file="BilateralFilter.cs" company="Microsoft Corporation">
// The MIT License (MIT)
// 
// Copyright (c) 2014, Microsoft Corporation
// 
// Permission is hereby granted, free of charge, to any person obtaining a copy
//  of this software and associated documentation files (the "Software"), to deal
//  in the Software without restriction, including without limitation the rights
//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//  copies of the Software, and to permit persons to whom the Software is
//  furnished to do so, subject to the following conditions:
// 
// The above copyright notice and this permission notice shall be included in
//  all copies or substantial portions of the Software.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
//  THE SOFTWARE.
// </copyright>
// --------------------------------------------------------------------------------------------------------------------

namespace Microsoft.Robotics.Vision
{
    using System;
    using System.Collections.Generic;
    using System.Linq;
    using System.Runtime.CompilerServices;
    using System.Threading.Tasks;
    using Microsoft.Robotics.Numerics;
    
    /// <summary>
    /// Implements a bilateral filter that can be applied to an entire depth image or to one pixel at a time
    /// </summary>
    public class BilateralFilter
    {
        /// <summary>
        /// This constant is the inverse depth (in millimeters) at which 'rangeSigma' is used without alteration
        /// </summary>
        private const double InverseNominalDepthForRangeSigma = 1 / 1000d;  

        /// <summary>
        /// The spatial gaussian kernel
        /// </summary>
        private double[,] spatialKernel;

        /// <summary>
        /// The size of the kernel
        /// </summary>
        private int kernelSize;
        
        /// <summary>
        /// Half size of the kernel
        /// </summary>
        private int halfKernelSize;

        /// <summary>
        /// The sigma for the range kernel
        /// </summary>
        private double rangeSigma;

        /// <summary>
        /// If true, the range sigma is interpreted as a range at 1m and adjusted accordingly
        /// </summary>
        private bool autoAdjustRangeSigma;

        /// <summary>
        /// Initializes a new instance of the BilateralFilter class
        /// </summary>
        /// <param name="spatialSigma">standard deviation of spatial Gaussian (in pixels, recommended value:  0.7)</param>
        /// <param name="rangeSigma">standard deviation of range Gaussian (in millimeters, recommended value:  10.0) </param>
        /// <param name="autoAdjustRangeSigma">if true, then the standard deviation of the range Gaussian is computed automatically, using 'rangeSigma' as the value for a depth of 1 meter</param>
        public BilateralFilter(double spatialSigma, double rangeSigma, bool autoAdjustRangeSigma)
        {
            this.spatialKernel = ImageOperations2D.Construct2DGaussianKernel(spatialSigma);
            this.kernelSize = this.spatialKernel.GetLength(0);
            this.halfKernelSize = (this.kernelSize - 1) / 2;
            this.rangeSigma = rangeSigma;
            this.autoAdjustRangeSigma = autoAdjustRangeSigma;
        }

        /// <summary>
        /// Performs bilateral filtering on a depth image using the standard, obvious, slow implementation
        /// </summary>
        /// <param name="image">Input depth image (assumed to be in millimeters)</param>
        /// <returns>The bilaterally filtered depth image</returns>
        public ImageFrameDepth Apply(ImageFrameDepth image)
        {
            ImageFrameDepth filtered = new ImageFrameDepth(image.Width, image.Height);

            Parallel.For(
                this.halfKernelSize,
                image.Height - this.halfKernelSize,
                y =>
                {
                    for (int x = this.halfKernelSize; x < image.Width - this.halfKernelSize; x++)
                    {
                        short depthAtPixel = this.Apply(image, y, x);

                        // Store normalized weighted sum into output pixel
                        filtered[x, y] = depthAtPixel;
                    }
                });

            return filtered;
        }

        /// <summary>
        /// Performs bilateral filtering on the specified pixel
        /// </summary>
        /// <param name="image">Input depth image (assumed to be in millimeters)</param>
        /// <param name="x">Horizontal image coordinate</param>
        /// <param name="y">Vertical image coordinate</param>
        /// <returns>The filtered pixel value</returns>
        public short Apply(ImageFrameDepth image, int x, int y)
        {
            short depthAtPixel = image[x, y];
            if (depthAtPixel > 0 && x >= this.halfKernelSize && x < image.Width - this.halfKernelSize && y >= this.halfKernelSize && y < image.Height - this.halfKernelSize)
            {
                // Compute sigma for range Gaussian by applying d^2 formula
                double rangeSigmaForThisPixel = this.rangeSigma;
                if (this.autoAdjustRangeSigma)
                {
                    rangeSigmaForThisPixel = AutoComputeRangeSigma(depthAtPixel, InverseNominalDepthForRangeSigma, this.rangeSigma);
                }

                double val = 0;
                double normalizer = 0;
                double inverseVariance = 1 / (2 * rangeSigmaForThisPixel * rangeSigmaForThisPixel);
                for (int dy = 0; dy < this.kernelSize; dy++)
                {
                    for (int dx = 0; dx < this.kernelSize; dx++)
                    {
                        short depthAtNeighbor = image[x + dx - this.halfKernelSize, y + dy - this.halfKernelSize];
                        if (depthAtNeighbor > 0)
                        {
                            // Compute weight for pixel by combining spatial and range weights
                            double difference = depthAtPixel - depthAtNeighbor;

                            // TODO: replace gaussian with table lookup, as this line is 24% of the function according to the profiler
                            double weight = this.spatialKernel[dx, dy] * Math.Exp(-difference * difference * inverseVariance);
                            normalizer += weight;

                            // Accumulate weighted sum
                            val += depthAtNeighbor * weight;
                        }
                    }
                }

                // Normalize weighted sum and clamp
                val = Math.Round(val / normalizer);
                if (val < 0)
                {
                    val = 0;
                }

                if (val > short.MaxValue)
                {
                    val = short.MaxValue;
                }

                depthAtPixel = (short)val;
            }

            return depthAtPixel;
        }

        /// <summary>
        /// Automatically compute the sigma for the range Gaussian for a particular depth
        /// </summary>
        /// <param name="depth">The depth of the pixel for which to compute the range sigma</param>
        /// <param name="inverseNominalDepth">The depth at which the returned value is equal to the nominal range sigma</param>
        /// <param name="nominalRangeSigma">The nominal range sigma</param>
        /// <returns>The range sigma automatically adjusted for depth</returns>
        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        internal static double AutoComputeRangeSigma(double depth, double inverseNominalDepth, double nominalRangeSigma)
        {
            double relativeDepth = depth * inverseNominalDepth - 1;
            return nominalRangeSigma * (1 + relativeDepth * relativeDepth);
        }
    }
}
